import os

from langchain_community.chat_models.zhipuai import ChatZhipuAI
from langchain_core.messages import SystemMessage, RemoveMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import MessagesState, StateGraph, START, END

# os.environ["ZHIPUAI_API_KEY"] = "97738d4998b8732d707daf91a2b1c56d.2y6VKEuOlidwHDpI"

memory = MemorySaver()

class State(MessagesState):
    summary: str

model = ChatZhipuAI(
    model = "glm-4",
    temperature=0.95,
)

def call_model(state: State):
    summary = state.get("summary", "")
    if summary:
        system_message = f"Summary of conversation earlier: {summary}"
        messages = [SystemMessage(content=system_message)] + state["messages"]
    else:
        messages = state["messages"]
    response = model.invoke(messages)
    return {"messages": response}

def should_continue(state: State):
    messages = state["messages"]
    if len(messages) > 6:
        return "summarize_conversation"
    return END

def summarize_conversation(state: State):
    summary = state.get("summary", "")
    if summary:
        summary_message = (
            f"This is summary of the conversation to date: {summary}\n\n"
            "Extend the summary by taking into account the new messages above:"
        )
    else:
        summary_message = "Create a summary of the conversation above:"
    messages = state["messages"] + [HumanMessage(content=summary_message)]
    response = model.invoke(messages)
    return {"summary": response.content, "messages": [RemoveMessage(id=m.id) for m in state["messages"][:-2]]}

builder = StateGraph(State)
builder.add_node("call_model", call_model)
builder.add_node("summarize_conversation", summarize_conversation)

builder.add_edge(START, "call_model")
builder.add_conditional_edges(
    "call_model",
    should_continue,
    ["summarize_conversation", END]
)
builder.add_edge("summarize_conversation", END)
graph = builder.compile(checkpointer=memory)

# def stream_message(content: str):
#     config = {"configurable": {"thread_id": "4"}}
#     input_message = HumanMessage(content=content)
#     for chunk in graph.stream({"messages": [input_message]}, config, stream_mode="values"):
#         chunk["messages"][-1].pretty_print()
#     print(graph.get_state(config).values)
#
# stream_message("hi! I'm bob")
# stream_message("what's my name?")
# stream_message("i like the celtics!")
# stream_message("i like how much they win")
# stream_message("what's my name?")
# stream_message("what NFL team do you think I like?")
# stream_message("i like the patriots!")
